import logging
import time
from functools import cached_property
from typing import Any
import numpy as np

from lerobot.utils.errors import DeviceAlreadyConnectedError, DeviceNotConnectedError

from ..robot import Robot
from .config_so101_follower import SO101FollowerConfig

logger = logging.getLogger(__name__)


class SO101FollowerMock(Robot):
    """
    Mock SO-101 Follower Arm for testing without real hardware.
    All operations are simulated in memory.
    """

    config_class = SO101FollowerConfig
    name = "so101_follower_mock"

    def __init__(self, config: SO101FollowerConfig):
        super().__init__(config)
        self.config = config
        
        # Mock motor names
        self.motor_names = [
            "shoulder_pan",
            "shoulder_lift", 
            "elbow_flex",
            "wrist_flex",
            "wrist_roll",
            "gripper"
        ]
        
        # Mock state
        self._is_connected = False
        self._is_calibrated = True  # Mock is always "calibrated"
        
        # Initialize mock motor positions (middle of range)
        self._motor_positions = {motor: 0.0 for motor in self.motor_names}
        self._motor_positions["gripper"] = 50.0  # gripper at half-open
        
        self.cameras = {}
    
        # Mock camera data
        if config.cameras:
            for cam_name, cam_config in config.cameras.items():
                self.cameras[cam_name] = {
                    "height": cam_config.height,
                    "width": cam_config.width,
                }

    @property
    def _motors_ft(self) -> dict[str, type]:
        return {f"{motor}.pos": float for motor in self.motor_names}

    @property
    def _cameras_ft(self) -> dict[str, tuple]:
        return {
            cam: (info["height"], info["width"], 3) 
            for cam, info in self.cameras.items()
        }

    @cached_property
    def observation_features(self) -> dict[str, type | tuple]:
        return {**self._motors_ft, **self._cameras_ft}

    @cached_property
    def action_features(self) -> dict[str, type]:
        return self._motors_ft

    @property
    def is_connected(self) -> bool:
        return self._is_connected

    def connect(self, calibrate: bool = True) -> None:
        """Mock connection - just sets flag."""
        if self.is_connected:
            raise DeviceAlreadyConnectedError(f"{self} already connected")
        
        self._is_connected = True
        logger.info(f"{self} (MOCK) connected.")

    @property
    def is_calibrated(self) -> bool:
        return self._is_calibrated

    def calibrate(self) -> None:
        """Mock calibration - does nothing."""
        logger.info(f"{self} (MOCK) calibration skipped - mock is always calibrated.")
        self._is_calibrated = True

    def configure(self) -> None:
        """Mock configuration - does nothing."""
        logger.info(f"{self} (MOCK) configuration skipped.")

    def setup_motors(self) -> None:
        """Mock motor setup - does nothing."""
        logger.info(f"{self} (MOCK) motor setup skipped.")

    def get_observation(self) -> dict[str, Any]:
        """Return mock observation data."""
        if not self.is_connected:
            raise DeviceNotConnectedError(f"{self} is not connected.")

        obs_dict = {}
        
        # Mock motor positions
        for motor in self.motor_names:
            obs_dict[f"{motor}.pos"] = self._motor_positions[motor]
        
        # Mock camera images (random noise for now)
        for cam_name, cam_info in self.cameras.items():
            # Generate random RGB image
            mock_image = np.random.randint(
                0, 255, 
                (cam_info["height"], cam_info["width"], 3),
                dtype=np.uint8
            )
            obs_dict[cam_name] = mock_image
        
        return obs_dict

    def send_action(self, action: dict[str, Any]) -> dict[str, Any]:
        """Mock action send - updates internal state."""
        if not self.is_connected:
            raise DeviceNotConnectedError(f"{self} is not connected.")

        # Extract goal positions from action
        goal_pos = {}
        for key, val in action.items():
            if key.endswith(".pos"):
                motor_name = key.removesuffix(".pos")
                goal_pos[motor_name] = val
        
        # Apply max_relative_target constraint if configured
        if self.config.max_relative_target is not None:
            for motor, target in goal_pos.items():
                current = self._motor_positions[motor]
                diff = target - current
                if abs(diff) > self.config.max_relative_target:
                    # Clip the difference
                    diff = np.sign(diff) * self.config.max_relative_target
                    goal_pos[motor] = current + diff
        
        # Update mock positions
        for motor, pos in goal_pos.items():
            self._motor_positions[motor] = pos
        
        # Return the actual commanded positions
        return {f"{motor}.pos": val for motor, val in goal_pos.items()}

    def disconnect(self):
        """Mock disconnection - just clears flag."""
        if not self.is_connected:
            raise DeviceNotConnectedError(f"{self} is not connected.")

        self._is_connected = False
        logger.info(f"{self} (MOCK) disconnected.")